import torch.nn as nn
from torch.nn import Sequential

class VGGNet(nn.Module):
    # 定义配置信息
    cfg = {
        # vgg11网络配置
        "vgg11": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
        "vgg13": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
        "vgg16": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"],
        "vgg19": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"]
    }

    def __init__(self, net_type, num_classes):
        """
        初始化配置
        :param net_type: 调用的网络类型
        """
        super().__init__()
        # 生成网络
        self.__create_nn(net_type, num_classes)

    def __create_nn(self, net_type, num_classes):
        """
        生成网络
        :param net_type: 网络类型
        :return:
        """
        in_channel = 3
        base_config = []
        for item in self.cfg[net_type]:
            if item == "M":
                base_config.append(nn.MaxPool2d((2, 2), (2, 2)))
            else:
                base_config.append(nn.Conv2d(in_channel, item, (3, 3), (1, 1), (1, 1)))
                base_config.append(nn.ReLU(True))
                in_channel = item

        # 分类网络
        classifier = self.__classifier(num_classes)
        base_config.extend(classifier)
        # 注册网络
        self.sequential = Sequential(*base_config)
        # 初始化参数
        self.__load_args()

    def __classifier(self, num_classes):
        classifier = [
            nn.Flatten(),
            nn.Dropout(0.5),
            nn.Linear(512*7*7, 4096),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(4096, num_classes)
        ]
        return classifier

    def __load_args(self):
        # 初始化权重参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
                nn.init.constant_(m.bias, 0.0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.constant_(m.bias, 0.0)

    def forward(self, x):
        r = self.sequential(x)
        return r

def vgg16(classes):
    return VGGNet("vgg16", classes)

def vgg19(classes):
    return VGGNet("vgg19", classes)

if __name__ == "__main__":
    from utils.utils import get_model_flops_args
    get_model_flops_args(VGGNet("vgg16", 5), (3, 224, 224))